import os
import argparse
import sys

samtools='/data/tools/samtools-1.6/samtools'
gatk='/data/zhanglk/biosoft/GATK/gatk-4.0.2.1/gatk'
bwa='/data/tools/bwa-0.7.17/bwa'

parser = argparse.ArgumentParser(description='输入参数如下:')
parser.add_argument('--process', '-P', type=str, help='进程数量，必要参数', required=True)
parser.add_argument('--fastq1', '-1', help='fastq1 file 必要参数', required=True)
parser.add_argument('--fastq2', '-2', help='fastq2 file 必要参数', required=True)
parser.add_argument('--genome', '-G', help='reference genome 必要参数', required=True)
parser.add_argument('--output', '-O', help='output fold 必要参数', required=True)
parser.add_argument('--prefix', '-N', help='prefix name 必要参数', required=True)
args = parser.parse_args()
fq1=args.fastq1 # fastq
fq2=args.fastq2 # fastq
if fq1==fq2:
    print('Both fastq files are the same!!')
    sys.exit()
genome=args.genome
if genome.endswith('.fa') or genome.endswith('.fasta'):
    print('Genome format right')
else:
    print('Genome format wrong, must end with .fa or .fasta')
    sys.exit()
out_prefix=args.prefix
out_fold=args.output
out_name=out_fold+'/'+out_prefix
thread=args.process
os.system('mkdir '+out_fold)
if os.path.exists(genome+'.bwt'):
    print('bwa index OK!')
else:
    os.system(bwa+' index '+genome)

if os.path.exists(genome+'.fai'):
    print('samtools faidx OK!')
else:
    os.system(samtools+' faidx '+genome)

if os.path.exists(genome[:-(len(genome.split('.')[-1])+1)]+'.dict'):
    print('gatk dict OK!')
else:
    dict_command= "%s CreateSequenceDictionary -R %s -O %s.dict"  %(gatk,genome,genome[:-(len(genome.split('.')[-1])+1)] )
    os.system(dict_command)
# sys.exit()

###### bwa 比对 ######
bwa_command="%s mem -R '@RG\\tID:foo\\tLB:library1\\tPU:runbarcode\\tPL:Illumina\\tDS:resequencing\\tCN:MajorBio\\tSM:%s' \
-t %s -M %s %s %s -o %s.sam" %(bwa,out_name,thread,genome,fq1,fq2,out_name)
print(bwa_command)
os.system(bwa_command)

###### samtools view ######
view_command="%s view -bS -h -@ %s %s.sam -o  %s.bam" %(samtools,thread,out_name,out_name)
os.system(view_command)

###### samtools sort ######
sort_command='%s sort -@ %s %s.bam -o %s.sorted.bam' %(samtools,thread,out_name,out_name)
os.system(sort_command)

###### gatk MarkDuplicates ######
dup_command="%s MarkDuplicates -I %s.sorted.bam -O %s.sorted.markdup.bam -M %s.sorted.markdup_metrics.txt" \
    %(gatk,out_name,out_name,out_name)
os.system(dup_command)

###### samtools index ######
index_command="%s index -@ %s %s.sorted.markdup.bam"  %(samtools,thread,out_name,)
os.system(index_command)

###### gatk HaplotypeCaller ######
gatk_command="%s  HaplotypeCaller --emit-ref-confidence GVCF -R %s -I %s.sorted.markdup.bam -O %s.g.vcf" %(gatk,genome,out_name,out_name)
os.system(gatk_command)